#coding=utf-8
from collections import defaultdict
import cv2
import os
import os.path as osp
import copy
import sys
import numpy as np
import shutil
import time, json
from multiprocessing import Pool
from datetime import datetime
from PIL import ImageFont, ImageDraw, Image

# from tool import video_tool # export PYTHONPATH=$PYTHONPATH:`pwd`
sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))
from tool import filesystem


def capture():
    # path = sys.argv[1]
    # for p in os.listdir(path):
    #     video_to_img(path + os.sep + p)

    cap = cv2.VideoCapture(0)

    # cap = cv2.VideoCapture(1)
    width = 1280
    height = 720
    cap.set(3,width)   # width
    cap.set(4,height)   # height
    cap.set(6, cv2.VideoWriter.fourcc("M","J","P","G"))
    # cap.set(cv2.CAP_PROP_EXPOSURE, 1000)


    fourcc = cv2.VideoWriter_fourcc("m","p","4","v")
    start_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))

    output_movie = cv2.VideoWriter('{}.mp4'.format(start_time), fourcc, 15, (int(width), int(height)))

    count = 1
    save_video = False
    while 1:
        # if count % 10 == 0:
        #     cap.set(10, count / 5)   # height


        ret, frame = cap.read()
        # frame = cv2.resize(frame, (1920, 1080))
        cv2.imshow("capture",frame)
        #cv2.moveWindow('capture',50,50)

        if save_video:
            # output_movie.write(frame)
            start_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))

            cv2.imwrite("{}.png".format(start_time), frame)
        # if count % 3 == 0:
        #     cv2.imwrite('./pics/{}.jpg'.format(str(count)), frame)
            
        if cv2.waitKey(1) & 0xFF == ord('b'):
            break
        if cv2.waitKey(1) & 0xFF == ord('s'):
            print("start to save video...")
            save_video = True
        count += 1

    cap.release()
    cv2.destroyAllWindows()
    output_movie.release()

def video_to_img(input_path, dest_dir=None, step=3):
    """
    dest_dir: str
        将视频转成图片 并保存到对应目录下 %未实现%
    step: int
        跳帧
    """
    ## get to list
    if os.path.isdir(input_path):
        video_paths = filesystem.get_all_filepath(input_path, [".mp4", ".avi", ".mov"])
    else:
        video_paths = [input_path]
    print(video_paths)
    
    ## save image
    for video_path in video_paths:

        base_name  = os.path.basename(video_path)
        new_dir = os.path.dirname(video_path) + os.sep + ".".join(base_name.split('.')[:-1])
        if not os.path.exists(new_dir):
            os.makedirs(new_dir, exist_ok=True)

        cap = cv2.VideoCapture(video_path)
        fps = cap.get(5)
        if fps< 1 or fps > 100:
            fps = 25.
            cap.set(5, fps)

        print(fps)

        count = 0
        while 1:
            ret, frame = cap.read()
            if not ret:
                print("cap.read() {} is over...".format(base_name))
                break

            # if count < 45223:continue
            # cv2.imshow("capture",frame)
            # cv2.moveWindow('capture',100,100)
            
            if count % step == 0:
                # frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
                cv2.imwrite(new_dir + os.sep +'{}_{}.jpg'.format(base_name, count), frame) # [:, 420:1500, :]
            # print(np.mean(frame))
            print(count)
            count += 1
            
                
            if cv2.waitKey(1) & 0xFF == ord('b'):
                break
        cap.release()
        # cv2.destroyAllWindows()


def video_to_img_thread(input_path, dest_dir=None, step=3):
    """
    dest_dir: str
        将视频转成图片 并保存到对应目录下 %未实现%
    step: int
        跳帧
    """
    ## get to list
    if os.path.isdir(input_path):
        video_paths = filesystem.get_all_filepath(input_path, [".mp4", ".avi", ".MP4"])
    else:
        video_paths = [input_path]
    print(video_paths)
    
    ## save image
    thread_count  = 12
    p = Pool(thread_count)
    for video_path in video_paths:
        print(video_path)
        p.apply_async(video_to_img, args=(video_path, None, step))

    p.close()
    p.join()

def video_to_img_by_cnt(input_path, dest_dir=None, start=0, end=1e5, step=3):
    """
    dest_dir: str
        将视频转成图片 并保存到对应目录下 %未实现%
    step: int
        跳帧
    """
    ## get to list

    base_name  = os.path.basename(input_path)
    new_dir = osp.join(os.path.dirname(input_path),".".join(base_name.split('.')[:-1]), str(start))
    if not os.path.exists(new_dir):
        os.makedirs(new_dir, exist_ok=True)

    cap = cv2.VideoCapture(input_path)
    count = 0
    while 1:
        ret, frame = cap.read()
        count += 1
        if not ret:
            print("cap.read() {} is over...".format(base_name))
            break

        # if count < 45223:continue
        # cv2.imshow("capture",frame)
        # cv2.moveWindow('capture',100,100)
        if count < start: continue
        if count > end: break

        if count % step == 0:
            # frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
            cv2.imwrite(new_dir + os.sep +'{}_{}.jpg'.format(base_name, count), frame) # [:, 420:1500, :]
        # print(np.mean(frame))
        print(count)
        
            
        if cv2.waitKey(1) & 0xFF == ord('b'):
            break
    cap.release()
    # cv2.destroyAllWindows()


def video_to_img_by_pair_img_idx(img_dir, mp4_dir, step):
    """
    将成对的图片区间从视频中取出
    dest_dir: str
        
    step: int
        跳帧
    """
    img_paths = filesystem.get_all_filepath(img_dir, [".jpg"])

    img_idx_data = []
    for img_path in img_paths:
        base_name = osp.basename(img_path)
        strs = base_name.split("_")
        mp4_idx = int(strs[0].split(".")[0])
        img_idx = int(strs[1].split(".")[0])
        img_idx_data.append([img_path, mp4_idx, img_idx])
    
    img_idx_data = sorted(img_idx_data, key=lambda x : (x[1], x[2]) )
    img_idx_dict = defaultdict(list)
    for idx in range(0, len(img_idx_data), 2):
        img_idx_dict[img_idx_data[idx][1]].append([img_idx_data[idx][2], img_idx_data[idx+1][2]])
    print(img_idx_dict)

    for mp4_idx, img_idx_pairs in img_idx_dict.items():
        input_path = osp.join(mp4_dir, "{}.mp4".format(mp4_idx))
        base_name  = os.path.basename(input_path)
        new_dir = osp.join(os.path.dirname(input_path),".".join(base_name.split('.')[:-1]))
        if not os.path.exists(new_dir):
            os.makedirs(new_dir, exist_ok=True)

        img_idx_pair_idx = 0
        start_idx = img_idx_pairs[img_idx_pair_idx][0]
        end_idx = img_idx_pairs[img_idx_pair_idx][1]
        cap = cv2.VideoCapture(input_path)
        count = 0
        while 1:
            ret, frame = cap.read()
            count += 1
            if not ret:
                print("cap.read() {} is over...".format(base_name))
                break

            if count < start_idx: continue
            if count >= end_idx:
                img_idx_pair_idx+=1
                if img_idx_pair_idx < len(img_idx_pairs):
                    start_idx = img_idx_pairs[img_idx_pair_idx][0]
                    end_idx = img_idx_pairs[img_idx_pair_idx][1]
                else:
                    break
            
            if count % step != 0: continue
                # frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
            cv2.imwrite(new_dir + os.sep +'{}_{}.jpg'.format(base_name, count), frame) # [:, 420:1500, :]
            # print(np.mean(frame))
            print(count)
            
        cap.release()

def img_to_video(input_dir, fps=1):
    """
    dest_dir: str
        将视频转成图片 并保存到对应目录下 %未实现%
    step: int
        跳帧
    """
    img_files = filesystem.get_all_filepath(input_dir, [".jpg"])
    img_files = sorted(img_files, key=lambda x: osp.basename(x))
    # img_path = r"E:\work\dataset\video\lbld\512\0\512.mp4_{}.jpg"
    # img_files = [img_path.format(i) for i in range(1553, 1965) ]

    width = 1920
    height = 1080
    fourcc = cv2.VideoWriter_fourcc("M","J","P","G")
    output_movie = cv2.VideoWriter('{}.mp4'.format(input_dir), fourcc,25, (int(width), int(height)))

    for img_file in img_files:
        print(osp.basename(img_file))
        img = cv2.imread(img_file)
        img = cv2.resize(img, (width, height))

        # cv2.imshow("t", img)
        # cv2.waitKey(5)
        # for _ in range(int(30/fps)):
        output_movie.write(img)
            
    output_movie.release()


def merge_img_to_video_draw_box(input_dir, fps=1, via_name="via_region_data.json"):
    """
    dest_dir: str
        将视频转成图片 并保存到对应目录下 %未实现%
    step: int
        跳帧
    """

    with open(input_dir + os.sep + via_name,"r") as rf:
        data_dict = json.loads(rf.read())
    data_new =  {}
    for k, value in data_dict.items():
        filename = value["filename"]
        data_new[filename] = value

    img_files = filesystem.get_all_filepath(input_dir, [".jpg"])
    # osp.basename(x).split("_")[1]
    # osp.basename(x).split(".")[-2].split("_")[-1]
    img_files = sorted(img_files, key=lambda x: int(osp.basename(x).split(".")[-2].split("_")[-1]))
    # img_path = r"E:\work\dataset\video\lbld\509\deal\509.mp4_{}_Copy.jpg"
    # img_files = [img_path.format(i) for i in range(534, 745) ]
    h,w = cv2.imread(img_files[0]).shape[:2]
    width = 1280
    height = 720
    fps = 30.0
    fourcc = cv2.VideoWriter_fourcc("M","J","P","G")
    output_movie = cv2.VideoWriter('{}.avi'.format(input_dir), fourcc,fps, (int(width), int(height)))

    for img_file in img_files:
        print(osp.basename(img_file))
        img = cv2.imread(img_file)
        img = cv2.resize(img, (width, height))
        filename = osp.basename(img_file)
        for region in data_new[filename]["regions"]:
            name = region["shape_attributes"]["name"]
            label = region["region_attributes"]["label"]
            if name == "polygon":
                points = np.array([[x,y] for x,y in zip(
                                region['shape_attributes']["all_points_x"], 
                                region['shape_attributes']["all_points_y"])], 
                                dtype=np.int)
                bg = np.zeros(img.shape, dtype=np.uint8) + np.array([0,255,0], dtype=np.uint8)
                bg = cv2.addWeighted(img, 0.7, bg, 0.3, 0.)

                mask1 = np.zeros(img.shape[:2], dtype=np.uint8)
                mask1 = cv2.fillPoly(mask1, [points], [255])

                img = cv2.copyTo(bg, mask1, img)
            elif name == "rect":

                x0 = int(region["shape_attributes"]["x"] * width / w)
                y0 = int(region["shape_attributes"]["y"] * height / h)
                x1 = int((region["shape_attributes"]["width"] + region["shape_attributes"]["x"]) * width / w)
                y1 = int((region["shape_attributes"]["height"] + region["shape_attributes"]["y"])  * height / h)
                
                if y0 <= 30: 
                    y0 = 30

                label = region["region_attributes"]["label"]
                if label == "boat":
                    color = (20,255,20)
                    line = 1
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    # cv2.putText(img, label, (x0, y0 - 10), cv2.FONT_HERSHEY_SIMPLEX, .5, (0, 0, 255), 2)
                
                elif label == "ringelman":
                    color = (20,20,255)
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    cv2.putText(img, label, (x0, y0 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
                
                else:
                    color = (0,0,255)
                    line = 3
                    
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    cv2.putText(img, label, (x0, y0 - 10), cv2.FONT_HERSHEY_SIMPLEX, .5, (0, 0, 255), 2)
                
        cv2.imshow("t", img)
        cv2.waitKey(5)
        # for _ in range(int(30/fps)):
        output_movie.write(img)
            
    output_movie.release()

def draw_chinese(img, text, left, top, textColor=(0, 255, 0), textSize=20):

    if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img)
    fontText = ImageFont.truetype("font/simsun.ttc", textSize, encoding="utf-8")
    draw.text((left, top), text, textColor, font=fontText)
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)


def iou(rect1, rect2):
    # x,y,w,h
    y0 = np.max([rect1[1],rect2[1]])
    y1 = np.min([rect1[1] + rect1[3],rect2[1]+rect2[3]])

    x0 = np.max([rect1[0],rect2[0]])
    x1 = np.min([rect1[0] + rect1[2],rect2[0]+rect2[2]])
    ret = x1 > x0 and y1 > y0
    print(rect1, rect2, ret)

    return int(ret)

def merge_img_to_video_draw_box_2(data_dir, cn_label={}, fps=1):

    data_new =  {}
    for via in filesystem.get_all_filepath(input_dir, [".json"]):
        with open(via,"r") as rf:
            data_dict = json.loads(rf.read())
        for k, value in data_dict.items():
            filename = value["filename"]
            if data_new.get(filename, None) is None:
                data_new[filename] = value
            else:
                data_new[filename]["regions"].extend(value["regions"])
                
    img_files = filesystem.get_all_filepath(input_dir, [".jpg"])
    # osp.basename(x).split("_")[1]
    # osp.basename(x).split(".")[-2].split("_")[-1]
    img_files = sorted(img_files, key=lambda x: int(osp.basename(x).split(".")[-2].split("_")[-1]))
    # img_path = r"E:\work\dataset\video\lbld\509\deal\509.mp4_{}_Copy.jpg"
    # img_files = [img_path.format(i) for i in range(534, 745) ]
    h,w = cv2.imread(img_files[0]).shape[:2]
    fourcc = cv2.VideoWriter_fourcc("M","J","P","G")
    output_movie = cv2.VideoWriter('{}.avi'.format(input_dir), fourcc,fps, (int(w), int(h)))

    for img_file in img_files:
        print(osp.basename(img_file))
        img = cv2.imread(img_file)
        filename = osp.basename(img_file)
        for region in data_new[filename]["regions"]:
            name = region["shape_attributes"]["name"]
            label = region["region_attributes"]["label"]
            if name == "polygon":

                if label in ["sand1", "sand2"]:
                    color = (20,20,255)  # bgr
                elif label in ["sand_ok"]:
                    color = (20,255,20)  # bgr


                points = np.array([[x,y] for x,y in zip(
                                region['shape_attributes']["all_points_x"], 
                                region['shape_attributes']["all_points_y"])], 
                                dtype=np.int32)
                bg = np.zeros(img.shape, dtype=np.uint8) + np.array(color, dtype=np.uint8)
                bg = cv2.addWeighted(img, 0.8, bg, 0.2, 0.)

                mask1 = np.zeros(img.shape[:2], dtype=np.uint8)
                mask1 = cv2.fillPoly(mask1, [points], [255])
                
                img = cv2.copyTo(bg, mask1, img)


                # rect = cv2.boundingRect(points)
                points = sorted(points, key=lambda x : x[0] * x[1])
                rect = points[0]
                if cn_label[label] != "":
                    img = draw_chinese(img, cn_label[label], rect[0]-96, rect[1] - 20, color[::-1], 30)
                    
                # rect = points[-1]
                # if cn_label[label] != "":
                #     img = draw_chinese(img, cn_label[label], rect[0]+50, rect[1] + 10, color[::-1], 30)
                    
            elif name == "rect":
                # x0 = int(region["shape_attributes"]["x"] * width / w)
                # y0 = int(region["shape_attributes"]["y"] * height / h)
                # x1 = int((region["shape_attributes"]["width"] + region["shape_attributes"]["x"]) * width / w)
                # y1 = int((region["shape_attributes"]["height"] + region["shape_attributes"]["y"])  * height / h)
                
                x0 = int(region["shape_attributes"]["x"])
                y0 = int(region["shape_attributes"]["y"])
                x1 = int((region["shape_attributes"]["width"] + region["shape_attributes"]["x"]) )
                y1 = int((region["shape_attributes"]["height"] + region["shape_attributes"]["y"]) )
                
                if y0 <= 30: 
                    y0 = 30

                label = region["region_attributes"]["label"]
                if label == "dust":
                    color = (20,20,255)  # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)
                elif label in ["truck.red"]:
                    color = (20,20,255) # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)
                elif label in ["truck.red2"]:
                    color = (20,20,255) # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)

                elif label in ["truck", "truck0", "truck2.car"]:
                    color = (255,20,20) # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)
                elif label in ["truck1", "truck2", "truck.dust"]:
                    alarm = 0
                    for region00 in data_new[filename]["regions"]:
                        if region00["region_attributes"]["label"] == "dust":
                            if iou((x0,y0,x1-x0,y1-y0), (int(region00["shape_attributes"]["x"]),int(region00["shape_attributes"]["y"]),int(region00["shape_attributes"]["width"]),int(region00["shape_attributes"]["height"]))) == 1:
                                alarm = 1
                    print("alarm:",alarm)
                    if alarm == 1:
                        color = (20,20,255) # bgr
                        text = cn_label[label]
                    else:
                        color = (255,20,20) # bgr
                        text = "大货车"
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, text, x0, y0 - 35, color[::-1], 30)

                elif label == "truck_clear":
                    color = (20,255,20) # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)

                elif label == "person":
                    color = (20,255, 20) # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)
                elif label == "person_nohat":
                    color = (20,20,255)     # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)
                elif label == "rolling_door_open":
                    color = (255,20,20)     # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)
                elif label == "rolling_door_closed":
                    color = (20,255,20)     # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)

                elif label in ["wupao", "truck3", "penlin", "sashui", ] :
                    color = (20, 255,20)     # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)
                elif label == "wupaoji":
                    color = (20,20, 255)     # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)
                elif label in ["wupaoji_work", "truck4.wupao"]:
                    color = (20, 255,20)     # bgr
                    line = 3
                    img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                    if cn_label[label] != "":
                        img = draw_chinese(img, cn_label[label], x0, y0 - 35, color[::-1], 30)

                # else:
                #     color = (0,0,255)
                #     line = 3
                #     img = cv2.rectangle(img, (x0,y0), (x1,y1), color, line)
                #     cv2.putText(img, label, (x0, y0 - 10), cv2.FONT_HERSHEY_SIMPLEX, .5, (0, 0, 255), 2)
                
        # cv2.imshow("t", img)
        # cv2.waitKey(5)
        # for _ in range(int(30/fps)):
        output_movie.write(img)
            
    output_movie.release()

def video_to_img_expand(video__dir, save_dir, step=None, expand_len=8):
    """
    根据已经存在的图片进行相应首尾补充
    save_dir: str
        将视频转成图片 并保存到对应目录下
    step: int
        跳帧
    """
    ## get to list
    if os.path.isdir(video__dir):
        video_paths = filesystem.get_all_filepath(video__dir, [".mp4", ".avi"])
    else:
        video_paths = [video__dir]
        video__dir = os.path.dirname(video__dir)
    
    def get_top_k(l, k):
        d = defaultdict(int)
        for i in l: d[i] += 1
        d = list(d.items())
        d2 = sorted(d, key=lambda x: x[1], reverse=True)
        d3 = [x[0] for x in d2]
        return d3[:k]

    ## save image
    for video_path in video_paths:

        base_name  = os.path.basename(video_path)
        new_dir = os.path.dirname(video_path).replace(video__dir, save_dir) + os.sep + ".".join(base_name.split('.')[:-1])
        # new_dir = ".".join(video_path.split('.')[:-1]).replace(video__dir, save_dir)
        if not os.path.exists(new_dir): 
            print("CONTINUE... ", video_path)
            continue
        
        exist_files = filesystem.get_all_filepath(new_dir, [".jpg"])
        exist_idx = [int(".".join(osp.basename(f).split(".")[:-1]).split("_")[-1]) for f in exist_files]
        exist_idx = sorted(exist_idx)
        exist_diff = [exist_idx[i]-exist_idx[i-1] for i in range(1, len(exist_idx))]
        step = get_top_k(exist_diff, 1)
        step = step[0]

        extract_frame_idx = [i for idx in exist_idx for i in range(idx- step*expand_len, idx+ step*expand_len, step)]
        extract_frame_idx = list(set(extract_frame_idx))

        cap = cv2.VideoCapture(video_path)
        count = 0
        while 1:
            ret, frame = cap.read()
            count += 1
            # cv2.imshow("capture",frame)
            # cv2.moveWindow('capture',100,100)
            if not ret:
                print("cap.read() {} is over...".format(base_name))
                break

            # if count < 45223:continue

            # if cv2.waitKey(1) & 0xFF == ord('c'):
            if count in extract_frame_idx:
                # frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
                save_path = new_dir + os.sep +'{}_{}.jpg'.format(base_name, count)
                if not osp.exists(save_path):
                    cv2.imwrite(save_path, frame) # [:, 420:1500, :]
            # print(np.mean(frame))
            print(count)
            
                
            # if cv2.waitKey(1) & 0xFF == ord('b'):
                # break
        cap.release()
        # cv2.destroyAllWindows()


def video_to_video(file_path):

    cap = cv2.VideoCapture(file_path)
    width = cap.get(3)
    height = cap.get(4)
    print(width, " " ,height)
    fourcc = cv2.VideoWriter_fourcc("M","J","P","G")
    output_movie = cv2.VideoWriter('11output.avi', fourcc, 22, (int(width), int(height)))


    # cap.set(6, cv2.VideoWriter.fourcc("M","J","P","G"))

    start_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))

    save_flag=False
    while 1:
        ret, frame = cap.read()

        if not ret:
            break
        cv2.imshow("t", frame)
        if save_flag:
            output_movie.write(frame)

        if cv2.waitKey(1) & 0xFF == ord('s'):
            save_flag = True
            print("start save...")
            
        if cv2.waitKey(1) & 0xFF == ord('b'):
            break
    cap.release()
    output_movie.release()

def video_crop_to_video(video_path, idx_pair):
    save_dir = osp.join(osp.dirname(video_path), ".".join(osp.basename(video_path).split(".")[:-1]))
    os.makedirs(save_dir, exist_ok=True)

    cap = cv2.VideoCapture(video_path)
    width = cap.get(3)
    height = cap.get(4)
    print(width, " " ,height)
    fourcc = cv2.VideoWriter_fourcc("M","J","P","G")
    count = 0

    for pair in idx_pair:
        save_path = osp.join(save_dir, "{}_{}.avi".format(pair[0], pair[1]))
        output_movie = cv2.VideoWriter(save_path, fourcc, 22, (int(width), int(height)))
        # cap.set(6, cv2.VideoWriter.fourcc("M","J","P","G"))

        while 1:
            ret, frame = cap.read()
            if not ret: break
            
            if count == pair[1]: break

            if pair[0] <= count < pair[1]:
                output_movie.write(frame)

            count+=1

        output_movie.release()
    cap.release()


def video_to_resize_video(video_path, dst_size):
    """
    生成不同大小的视频
    """
    cap = cv2.VideoCapture(video_path)
    width = cap.get(3)
    height = cap.get(4)
    fps = cap.get(5)
    print(width, " " ,height)
    fourcc = cv2.VideoWriter_fourcc("M","J","P","G")
    count = 0

    save_path = osp.join(osp.dirname(video_path), "{}_{}x{}.avi".format(".".join(osp.basename(video_path).split(".")[:-1]), dst_size[0], dst_size[1]))
    output_movie = cv2.VideoWriter(save_path, fourcc, fps, (int(dst_size[0]), int(dst_size[1])))
    # cap.set(6, cv2.VideoWriter.fourcc("M","J","P","G"))

    while 1:
        ret, frame = cap.read()
        if not ret: break
        
        new_frame = cv2.resize(frame, dst_size)
        output_movie.write(new_frame)

        count+=1

    output_movie.release()
    cap.release()


def video_to_video_2(file_path):
    """
    将横向视频转为竖直视频
    """
    cap = cv2.VideoCapture(file_path)
    width = cap.get(3)
    height = cap.get(4)
    print(width, " " ,height)
    fourcc = cv2.VideoWriter_fourcc("M","J","P","G")
    output_movie = cv2.VideoWriter('11output.avi', fourcc, 30, (int(height), int(width)))


    # cap.set(6, cv2.VideoWriter.fourcc("M","J","P","G"))

    start_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))

    save_flag=True
    while 1:
        ret, frame = cap.read()

        if not ret:
            break
        frame = cv2.rotate(frame, cv2.ROTATE_90_COUNTERCLOCKWISE)
        cv2.imshow("t", frame)
        if save_flag:
            output_movie.write(frame)

        if cv2.waitKey(1) & 0xFF == ord('s'):
            save_flag = True
            print("start save...")
            
        if cv2.waitKey(1) & 0xFF == ord('b'):
            break
    cap.release()
    output_movie.release()


def write_video():

    length = 10

    # Create an output movie file (make sure resolution/frame rate matches input video!)
    #fourcc = cv2.VideoWriter_fourcc(*'MP4V')
    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    output_movie = cv2.VideoWriter('00output2_9.mp4', fourcc, length,(1280, 720))

    frame_number = 0

    for i in range(30100, 45000):
        # Grab a single frame of video
        frame = cv2.imread("/media/swls/disk1/datasets/distance_recog/project_dataset/{}.jpg".format(str(i)))
        frame_number += 1
        # Write the resulting image to the output video file
        # print("Writing frame {} / {}".format(frame_number, length))
        output_movie.write(frame)

        if i % 1000 == 0:
            print(i)

    # All done!
    # cv2.destroyAllWindows()

# 递归取一个文件夹下所有文件到目标文件夹
count = 0
def move_all_to_one(all_dir_path, dest_dir):
    global count
    for f in os.listdir(all_dir_path):
        new_path = all_dir_path + os.sep + f
        if os.path.isfile(new_path) and f.endswith(".jpg"):
            dest_path = dest_dir + os.sep + str(count) + '_'  + os.path.basename(new_path)
            print(dest_path)
            count += 1
            shutil.move(new_path, dest_path)
        elif os.path.isdir(new_path):
            move_all_to_one(new_path, dest_dir)

def convert_to_jpg(input_path, dest_dir=None):
    file_paths = filesystem.get_all_filepath(input_path, [".bmp"])

    for file_path in file_paths:
        image = cv2.imread(file_path)
        save_path = file_path.replace(".bmp", ".jpg")
        cv2.imwrite(save_path, image)

def crop_video_by_key(video_path, interval=120):
    """
    interval: 秒
    """
    cnt = 0
    cap = cv2.VideoCapture(video_path)
    width = cap.get(3)
    height = cap.get(4)
    # fps = cap.get(cv2.CAP_PROP_FPS)
    fps = 20
    print(width, " " ,height, " fps: ", fps)

    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    output_movie = None

    start = 0
    start_time = time.time()
    while True:
        ret,frame = cap.read()
        if not ret: break
        
        frame_copy = copy.copy(frame)
        tmp_img = cv2.resize(frame_copy, (1920, 1080))
        cv2.imshow("w",tmp_img)
        key = cv2.waitKey(20)
        if key & 0xFF == ord('s'): 
            start = 1
            name = video_path + "_{}.crop.mp4".format(cnt)
            output_movie = cv2.VideoWriter(name, fourcc, fps, (int(frame.shape[1]), int(frame.shape[0])))
            cnt+=1
        if key & 0xFF == ord('e'): 
            start = 0
            output_movie.release()

        if key & 0xFF == ord('b'): 
            break

        if start: output_movie.write(frame)

        cv2.waitKey(10)
           
    cap.release()
    cv2.destroyWindow("w")

def split_video_to_sub_video(video_path, interval=120):
    """
    interval: 秒
    """
    cap = cv2.VideoCapture(video_path)
    width = cap.get(3)
    height = cap.get(4)
    fps = cap.get(cv2.CAP_PROP_FPS)
    print(width, " " ,height, " fps: ", fps)
    if fps > 100:
        fps = 15
    save_dir = ".".join(video_path.split(".")[:-1])
    os.makedirs(save_dir, exist_ok=True)

    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    writer = None
    # cap.set(6, cv2.VideoWriter.fourcc("M","J","P","G"))

    start_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
    video_cnt = 0
    count = 0
    while 1:
        ret, frame = cap.read()
        count+=1
        if not ret: break
        
        if count == 1:
            save_path = osp.join(save_dir, "{}.mp4".format(video_cnt))
            print(save_path)
            writer = cv2.VideoWriter(save_path, fourcc, fps, (int(width), int(height)))
            video_cnt+=1
        elif count > interval * fps :
            count = 0
            writer.release()
        else:
            writer.write(frame)


    cap.release()


def split_video_to_sub_video2(video_path, split_times):
    """
    interval: 秒
    """
    cap = cv2.VideoCapture(video_path)
    width = cap.get(3)
    height = cap.get(4)
    fps = cap.get(cv2.CAP_PROP_FPS)
    print(width, " " ,height, " fps: ", fps)
    if fps > 100:
        fps = 15
    save_dir = ".".join(video_path.split(".")[:-1])
    os.makedirs(save_dir, exist_ok=True)

    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    writer = None
    # cap.set(6, cv2.VideoWriter.fourcc("M","J","P","G"))

    split_idx = 0
    cur_s, cur_e = split_times[split_idx][0:2]
    is_write = 0
    count = 0
    while 1:
        ret, frame = cap.read()
        count+=1
        if not ret: break
        
        if cur_s * fps <= count <= cur_e * fps:
            is_write = 1

            if writer is None:
                save_path = osp.join(save_dir, "{}.mp4".format(count))
                print(save_path)
                writer = cv2.VideoWriter(save_path, fourcc, fps, (int(width), int(height)))

            writer.write(frame)

        elif count > cur_e * fps and is_write:
            is_write = 0
            writer.release()
            writer = None
            
            split_idx += 1
            if split_idx < len(split_times):
                cur_s, cur_e = split_times[split_idx][0:2]
            else:
                break
    cap.release()

def save_video_by_rtsp(save_dir, cap_ip, user, passwd, save_time):
    cap = cv2.VideoCapture("rtsp://{}:{}@{}:554/ch1/main/av_stream".format(user, passwd, cap_ip))
    print("ok")
    ret,frame = cap.read()
    print("ret: ", ret)
    if not ret: 
        return ret

    os.makedirs(save_dir, exist_ok=True)
    height, width = frame.shape[:2]
    fps = cap.get(5)
    print("fps: ", fps)
    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    name = '{}_{}.mp4'.format(cap_ip, datetime.now().strftime("%Y-%m-%d %H-%M-%S"))
    output_movie = cv2.VideoWriter(osp.join(save_dir, name), fourcc, 30, (int(width), int(height)))
    
    start_time = time.time()
    while True:
        ret,frame = cap.read()
        # cv2.imshow(cap_ip,frame)
        # cv2.waitKey(1)
        output_movie.write(frame)
        # if cv2.waitKey(1) & 0xFF == ord('q'): break

        if time.time() - start_time > save_time: break
           
    cap.release()
    output_movie.release()
    # cv2.destroyWindow(cap_ip)

def merge_video_to_one_video(video_dir):

    video_files = filesystem.get_all_filepath(video_dir, [".mp4", ".avi"])
    width = 0
    height = 0
    fourcc = cv2.VideoWriter_fourcc("M","J","P","G")
    save_path = osp.join(osp.dirname(video_dir), "{}.avi".format(osp.basename(video_dir)))
    output_movie = None

    for video_file in video_files:

        cap = cv2.VideoCapture(video_file)
        w = int(cap.get(3))
        h = int(cap.get(4))
        print(w, " " ,h)
        
        if width == 0 and height == 0:
            width = w
            height = h
        if output_movie is None:
            output_movie = cv2.VideoWriter(save_path, fourcc, 25, (int(width), int(height)))
            # cap.set(6, cv2.VideoWriter.fourcc("M","J","P","G"))

        while 1:
            ret, frame = cap.read()
            if not ret: break
            
            if frame.shape[0] != height or frame.shape[1] != width:
                frame = cv2.resize(frame, (width, height))
            output_movie.write(frame)


        cap.release()
    output_movie.release()


if __name__ == '__main__':

    # capture()

    video_to_img(sys.argv[1], None, step=1)

    # video_to_img(sys.argv[1], None, 1)
    # dir_path = r'E:\work\dataset\det\waterlogging\2023\07'
    # video_to_img(input_path=dir_path, dest_dir=None, step=3)
    # print(sys.argv[1])
    
    # video_to_img(sys.argv[1], None, 1)
    # video_to_img_thread(sys.argv[1], None, 30)

    # merge_video_to_one_video(sys.argv[1])

    # video_to_img_by_cnt(sys.argv[1], None, 6000, 6750, 1)

    # # 将手动筛选出的成对开始-结束图片区间从视频中提取出来
    # img_dir = r'F:\work\dataset\det\water\2025\0319\pic'
    # mp4_dir = r'F:\work\dataset\det\water\2025\0319\mp4'
    # video_to_img_by_pair_img_idx(img_dir, mp4_dir, step=10)
    
    # data = [
    #     [r"E:\work\dataset\video\lbld\11\610.mp4", 972, 1972],
    #     [r"E:\work\dataset\video\lbld\11\611.mp4", 4146, 5316],
    #     [r"E:\work\dataset\video\lbld\11\611.mp4", 7296, 8296],
    #     [r"E:\work\dataset\video\lbld\11\612.mp4", 5430, 6430],
    #     [r"E:\work\dataset\video\lbld\11\613.mp4", 9200, 10200],
    #     [r"E:\work\dataset\video\lbld\11\614.mp4", 9000, 11000]
    # ]
    # for input_path, s,e in data:
    #     video_to_img_by_cnt(input_path, None, s, e, 1)


    # video_path = sys.argv[1]
    # # 按时间等长分割
    # # split_video_to_sub_video(video_path, 100)
    # # 按指定时长分割
    # split_times = [[77, 120]]
    # split_video_to_sub_video2(video_path, split_times)

    # video_path = sys.argv[1]
    # crop_video_by_key(video_path)

    # save_dir = "/home/xc/work/code/ellipse_detector/video/2022-01-24"
    # cap_ip = "192.168.200.64"
    # user = "admin"
    # passwd = "abcd1234" 
    # save_time = 5*60
    # save_video_by_rtsp(save_dir, cap_ip, user, passwd, save_time)

    # idx_pair = [
    #         [3970, 4210],
    #         [5720, 5920],
    #         [5980, 6300],
    #         [7630, 7880],
    #         [8820, 8990],
    #         [10580, 10880],
    #     ]
    # video_crop_to_video(sys.argv[1], idx_pair)
    # video_to_resize_video(sys.argv[1], [640, 360])

    # video__dir = "/home/xc/work/data/car/2021-05-27"
    # save_dir = "/home/xc/work/code/paddle/train_data/det/car/images/2021-05-27"
    # video_to_img_expand(video__dir, save_dir, None, 64)

    # convert_to_jpg(sys.argv[1])
    # video_to_video_2(sys.argv[1])

    # video_to_video(r"C:\Users\swls\Kazam_screencast_00000.mp4")

    # input_dir = r"E:\work\dataset\test\598"
    # img_to_video(input_dir)
    # input_dir = r"E:\work\dataset\video\boat\pexels-peggy-anke-8787749"
    # merge_img_to_video_draw_box(input_dir)

    # label = {
    #     "person_nohat": "未戴安全帽",
    #     "person": "正常",
    #     "dust": "扬尘报警",
    #     "truck.dust": "大货车, 扬尘报警",
    #     "truck0": "大货车",
    #     "truck": "工作车",
    #     "truck.red": "工作车异常, 雾炮机未工作",
    #     "truck.red2": "工作车异常, 附近没有雾炮机",
    #     "truck1": "大货车",
    #     "truck2": "车辆",
    #     "truck2.car": "小车",
    #     "truck3": "正常作业",
    #     "truck4.wupao": "雾炮车正常作业",
    #     "truck_clear": "车辆正在清洗",
    #     "rolling_door_open": "正在开门",
    #     "rolling_door_close": "正在关门",
    #     "rolling_door_closed": "大门关闭",
    #     "wupaoji": "雾炮机未工作",
    #     "wupaoji_work": "雾炮机正在工作",
    #     "wupao": "雾炮",
    #     "penlin": "喷淋",
    #     "sashui": "洒水",

    #     "sand1": "未遮布",
    #     "sand2": "未遮布",
    #     "sand.green": "未遮布",
    #     "sand_ok": "已加湿",

    # }
    # input_dir = r"E:\work\dataset\det\sucai\4\1\0.ok"
    # fps = 30.0
    # merge_img_to_video_draw_box_2(input_dir, cn_label=label,fps=fps)


    # data_dir = sys.argv[1]
    # for f in os.listdir(data_dir):
    #     print(f)
    #     path= data_dir + os.sep + f
    #     if os.path.isfile(path) and f.endswith(".mp4"):
    #         video_to_img(path)

    # src = r"E:\work\data\electric_meter_box\video"
    # dest = r"E:\work\data\electric_meter_box\ic"
    # move_all_to_one(src, dest)
    
